# -*- coding: UTF-8 -*-
import torch
import torch.optim
import os
import numpy as np
from torchvision import transforms

from module.model import enhance_net_nopool

unloader = transforms.ToPILImage()
DCE_net = enhance_net_nopool().cuda()
DCE_net.load_state_dict(torch.load('/usr/projects/DVINS/module/snapshots/Epoch99.pth'))

def tensor_to_PIL(tensor):
    image = tensor.cpu().clone()
    image = image.squeeze(0)
    image = unloader(image)
    return image


def lowlight(frame):
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    # data_lowlight = Image.open(image_path)
    data_lowlight = (np.asarray(frame) / 255.0)
    data_lowlight = torch.from_numpy(data_lowlight).float()
    data_lowlight = data_lowlight.permute(2, 0, 1)
    data_lowlight = data_lowlight.cuda().unsqueeze(0)
    _, enhanced_image, _ = DCE_net(data_lowlight)
    return tensor_to_PIL(enhanced_image)


